import pickle
from model import multi_model
from utils.data_builder import data_builder
import time

from sklearn.metrics import accuracy_score, recall_score, f1_score, precision_score


import numpy as np

EPOCHES = 10
X_PATH = "data/sentence_codes_4096_dm0.npy" # 预训练模型
ORIG_PATH = "data/train.csv"
LABLE_NAME = ["toxic","severe_toxic","obscene","threat","insult","identity_hate"]

def timer(func):
    def _warp(*args, **kwargs):
        """
        :param args: func需要的位置参数
        :param kwargs: func需要的关键字参数
        :return: 函数的执行结果
        """
        start_time = time.time()
        result = func(*args, **kwargs)
        elastic_time = time.time() - start_time
        print("The execution time of the function '%s' is %.6fs" % (
            func.__name__, elastic_time))
        return result

    return _warp

def build_model_list():
    model_list = []
    for name in LABLE_NAME:
        f = open('out/{}.pickle'.format(name),'rb') 
        s = f.read()
        model_list.append(pickle.loads(s))
    return model_list



def validation(test_X, test_y, clf):
    y_pred = clf.predict_proba(test_X)
    y_round_pred = np.clip(np.around(y_pred).astype("int64"), 0, 1)

    print("predict_proba")
    print("accuracy_score", accuracy_score(test_y, y_round_pred))
    print("recall_score", recall_score(test_y, y_round_pred, average='micro'))
    print("f1_score", f1_score(test_y, y_round_pred, average='micro'))
    print("precision_score", precision_score(test_y, y_round_pred, average='micro'))
    

    y_pred = clf.predict_without_GNN(test_X)
    y_round_pred = np.clip(np.around(y_pred).astype("int64"), 0, 1)

    print("predict_without_GNN")
    print("accuracy_score", accuracy_score(test_y, y_round_pred))
    print("recall_score", recall_score(test_y, y_round_pred, average='micro'))
    print("f1_score", f1_score(test_y, y_round_pred, average='micro'))
    print("precision_score", precision_score(test_y, y_round_pred, average='micro'))
    

# 第二阶段，模型集成
# 载入全部的模型，对于一个test_X做预测, 得到若干个predict_y
# predict_y -> GNN -> 最终predict_y 
# 与真实标签进行比较，训练得到一个GNN集成学习器
def main():
    builder = data_builder(X_PATH, ORIG_PATH)
    train_loader, test_X, test_y = builder.build_graph_data()
    
    small_model_list = build_model_list()
    model = multi_model(small_model_list, EPOCHES)
    
    model.fit(train_loader)
    validation(test_X, test_y, model)

if __name__ == "__main__":
    main()